import os
from functools import partial
from multiprocessing.pool import Pool
from pathlib import Path

import cv2
import pandas as pd
from tqdm import tqdm

from constants import DEEPFAKE_DETECTION
from ffpp.preprocessing.generate_folds import cmp_frame_name, real_fake_split
from option import parse_args

os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"

cv2.ocl.setUseOpenCL(False)
cv2.setNumThreads(0)

SAMPLE_INTERVALS = [-40, -30, -20, -10, 0, 10, 20, 30, 40]


def get_real_fake_pairs(root_dir, original_path, manipulated_path, mode):
    lines = open(os.path.join(root_dir, '{}.txt'.format(mode)), 'r').readlines()
    actors = set([line.strip() for line in lines])
    pairs = []

    for video_name in os.listdir(os.path.join(original_path, 'videos')):
        video_name_split = video_name.split('__')
        if video_name_split[0] in actors:
            pair = (video_name[:-4], video_name[:-4])
            pairs.append(pair)

    for video_name in os.listdir(os.path.join(manipulated_path, 'videos')):
        video_name_split = video_name.split('__')
        assert len(video_name_split) == 3
        actor_ids = video_name_split[0]
        actor_ids_split = actor_ids.split('_')
        if actor_ids_split[0] in actors:
            pair = ('{}__{}'.format(actor_ids_split[0], video_name_split[1]), video_name[:-4])
            pairs.append(pair)
    return pairs


def check_crops2(ori_dir, fake_dir, frame):
    frame_split = frame.split('_')
    frame_num = int(frame_split[0])
    for interval in SAMPLE_INTERVALS:
        frame_cur = frame_num + interval
        img_name = f'{frame_cur}_{frame_split[1]}'
        # ori_img_path = os.path.join(ori_dir, img_name)
        fake_img_path = os.path.join(fake_dir, img_name)
        if not os.path.exists(fake_img_path):
            return False
    return True


def get_paths(vid, label, original_path, manipulated_path, sample_num: int, check_mask: bool):
    ori_vid, fake_vid = vid
    ori_dir = os.path.join(original_path, 'crops', ori_vid)
    if label == 1:
        fake_dir = os.path.join(manipulated_path, 'crops', fake_vid)
    else:
        fake_dir = os.path.join(original_path, 'crops', fake_vid)

    # mask: both c23 and c40 depend on c23 masks
    # fake_type_path = Path(manipulated_path).parent
    # manipulated_path_1 = os.path.join(fake_type_path, 'c23')
    # mask_dir = os.path.join(manipulated_path_1, 'diffs', fake_vid)

    data = []
    # Some unaligned videos have been removed
    if label == 1 and not os.path.exists(fake_dir):
        return data
    # Some masks may have been removed
    # if label == 1 and check_mask and not os.path.exists(mask_dir):
    #     return data

    frames = os.listdir(fake_dir)
    frames.sort(key=cmp_frame_name)
    for frame in frames:
        frame_split = str(frame).split('_')
        frame_num = int(frame_split[0])
        if frame_num % 50 != 0 or not check_crops2(ori_dir, fake_dir, frame):
            continue
        ori_img_path = os.path.join(ori_dir, frame)
        fake_img_path = os.path.join(fake_dir, frame)
        if label == 0 and os.path.exists(ori_img_path):
            data.append([ori_img_path, label, ori_vid])
        elif label == 1 and os.path.exists(fake_img_path):
            data.append([fake_img_path, label, ori_vid])
    return data


def save_folds(args, original_path, manipulated_path, save_path, pairs_mode, mode, sample_num=100, check_mask=True):
    data = []
    if mode == 'train_fake':
        ori_ori_pairs = set()
        ori_fake_pairs = pairs_mode
    elif mode == 'train_real':
        ori_ori_pairs = pairs_mode
        ori_fake_pairs = set()
    else:
        ori_ori_pairs = set([(ori, ori) for ori, fake in pairs_mode])
        ori_fake_pairs = set(pairs_mode) - ori_ori_pairs

    ori_ori_pairs = list(ori_ori_pairs)
    ori_fake_pairs = list(ori_fake_pairs)
    ori_ori_pairs.sort()
    ori_fake_pairs.sort()

    with Pool(processes=1) as p:
        # original label=0
        with tqdm(total=len(ori_ori_pairs)) as pbar:
            func = partial(get_paths, label=0, original_path=original_path, manipulated_path=manipulated_path,
                           sample_num=sample_num, check_mask=check_mask)
            for v in p.imap_unordered(func, ori_ori_pairs):
                pbar.update()
                data.extend(v)
        # fake label=1
        with tqdm(total=len(ori_fake_pairs)) as pbar:
            func = partial(get_paths, label=1, original_path=original_path, manipulated_path=manipulated_path,
                           sample_num=sample_num, check_mask=check_mask)
            for v in p.imap_unordered(func, ori_fake_pairs):
                pbar.update()
                data.extend(v)
    fold_data = []
    for img_path, label, ori_vid in data:
        path = Path(img_path)
        video = path.parent.name
        file = path.name
        fold_data.append([video, file, label, ori_vid, int(file.split("_")[0])])
    # random.shuffle(fold_data)
    columns = ["video", "file", "label", "original", "frame"]
    # save_file = '../data/{}/data_{}_{}_{}.csv'.format(FACE_FORENSICS, FACE_FORENSICS, args.fake_type, mode)
    file_name = '{}_st_{}.csv'.format(DEEPFAKE_DETECTION, mode)
    save_file = os.path.join(save_path, file_name)
    pd.DataFrame(fold_data, columns=columns).to_csv(save_file, index=False)


def main():
    args = parse_args()
    version = args.compression_version
    save_path = '../../data/{}/{}'.format(DEEPFAKE_DETECTION, version)
    os.makedirs(save_path, exist_ok=True)
    original_path = os.path.join(args.root_dir, 'original_sequences', 'actors', version)
    manipulated_path = os.path.join(args.root_dir, 'manipulated_sequences', 'DeepFakeDetection', version)
    pairs_train = get_real_fake_pairs(args.root_dir, original_path, manipulated_path, 'train')
    pairs_test = get_real_fake_pairs(args.root_dir, original_path, manipulated_path, 'test')
    pairs_val = get_real_fake_pairs(args.root_dir, original_path, manipulated_path, 'val')
    print(len(pairs_train), len(pairs_test), len(pairs_val))

    sample_num_train = 20
    sample_num_test = 20
    sample_num_val = 10
    check_mask = False
    save_folds(args, original_path, manipulated_path, save_path, pairs_train, mode='train',
               sample_num=sample_num_train, check_mask=check_mask)
    save_folds(args, original_path, manipulated_path, save_path, pairs_test, mode='test',
               sample_num=sample_num_test, check_mask=check_mask)
    save_folds(args, original_path, manipulated_path, save_path, pairs_val, mode='val',
               sample_num=sample_num_val, check_mask=check_mask)

    ori_fake_pairs_real, ori_fake_pairs_fake = real_fake_split(pairs_train)
    save_folds(args, original_path, manipulated_path, save_path, ori_fake_pairs_real, mode='train_real',
               sample_num=sample_num_train, check_mask=check_mask)
    save_folds(args, original_path, manipulated_path, save_path, ori_fake_pairs_fake, mode='train_fake',
               sample_num=sample_num_train, check_mask=check_mask)


if __name__ == '__main__':
    # Total: 2000
    # 'Deepfakes', 'Face2Face', 'FaceSwap', 'NeuralTextures', 'FaceShifter'
    main()
